load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRF_contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/contrib/ut_string.ncl"
load "/public/users/wrf/Codes/WRF_scripts/utils/substring.ncl"


undef("wrf_uv_intrp_vert")
function wrf_uv_intrp_vert(terrain:numeric, z:numeric, uvmet:numeric, uvmet10:numeric, y:numeric, x:numeric, height_list:numeric)
local hcol, ucol, vcol, u10, v10, alpha_u, alpha_v, alpha_u2, alpha_v2, uv_arr
begin

    dim_hlist = dimsizes(height_list)

    hcol = z(:, y, x) - terrain(y, x)
    ucol = uvmet(0, :, y, x)
    vcol = uvmet(1, :, y, x)
    u10 = uvmet10(0, y, x)
    v10 = uvmet10(1, y, x)

    ; compute alpha using u10 and ucol(1)
    if ( (ucol(1) * u10) .le. 0.0) then
    alpha_u = 0.143 ; default value of alpha
    else
    alpha_u = log(ucol(1)/u10)/log(hcol(1)/10.0)
    end if

    if ( (vcol(1) * v10) .le. 0.0) then
    alpha_v = 0.143 ; default value of alpha
    else
    alpha_v = log(vcol(1)/v10)/log(hcol(1)/10.0)
    end if

    ; compute alpha using ucol(0) and ucol(1)
    if ( (ucol(1) * ucol(0)) .le. 0.0) then
    alpha_u2 = 0.143 ; default value of alpha
    else
    alpha_u2 = log(ucol(1)/ucol(0))/log(hcol(1)/hcol(0))
    end if

    if ( (vcol(1) * vcol(0)) .le. 0.0) then
    alpha_v2 = 0.143 ; default value of alpha
    else
    alpha_v2 = log(vcol(1)/vcol(0))/log(hcol(1)/hcol(0))
    end if


    ; compute uv on height_list levels
;    u70 = u10 * (70.0/10.0)^alpha_u
;    v70 = v10 * (70.0/10.0)^alpha_v
;    u70_2 = ucol(0) * (70.0/hcol(0))^alpha_u2
;    v70_2 = vcol(0) * (70.0/hcol(0))^alpha_v2
    
    u_list = new(dim_hlist, "float")
    v_list = new(dim_hlist, "float")
    do ilvl = 0, dim_hlist - 1
     ;   print(ilvl + " level is: " + height_list(ilvl))
        htmp = height_list(ilvl)
        if (htmp .eq. 10) then
            utmp = u10
            vtmp = v10
        else
            utmp = u10 * (htmp/10.0)^alpha_u
            vtmp = v10 * (htmp/10.0)^alpha_v
        end if
        u_list(ilvl) = utmp
        v_list(ilvl) = vtmp
    end do

    uv_arr = (/u_list, v_list/)
    ;print(uv_arr)
    return(uv_arr)
end


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;            User Config                   ;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

; radius to degree
r2d = 45.0/atan(1.0)

if (.not. isvar("configFile")) then
    configFile = "/public/users/wrf/Codes/WRF_scripts/WRF_Post/StationLists/wind_stnlist.csv"
end if

lines = readAsciiTable(configFile,1,"string",1)
delim = ","

stnnum    = ndtooned(str_get_field(lines,1,delim))
stnlist    = ndtooned(str_get_field(lines,2,delim))
stnlon    = ndtooned(tofloat(str_get_field(lines,3,delim)))
stnlat    = ndtooned(tofloat(str_get_field(lines,4,delim)))
stndomain = ndtooned(str_get_field(lines,5,delim))
nstns     = dimsizes(stnlist)

levels     = (/10, 70/)
nlevels    = dimsizes(levels)

;;;;;;;;;;;;;;;;;;;;;;
; Env and Settings   ;
;;;;;;;;;;;;;;;;;;;;;;


nlevel = dimsizes(levels)

if (.not. isvar("output_prefix ")) then
    output_prefix = "~/wrfpost/test/station_wind"
    if (.not. isfilepresent(output_prefix)) then
        system("mkdir " + output_prefix)
    end if
end if

if (.not. isvar("input_file")) then
    input_file = "/public/users/wrf/wrfout/2012103120/wrfout_d02_2012-10-31_12:00:00"
end if

a = addfile(input_file + ".nc", "r")
filename = systemfunc("basename " + input_file)

if (.not. isvar("start_date")) then
    start_date = substring(filename, 11, 23)
end if

domain = substring(filename, 7, 9)

times = a->Times
hours = WRF_Times2Udunits_c(times, 0)
ntimes = dimsizes(times)

; UTC => CST
hours = hours + 8
timeStr_cst = ut_string(hours, "%Y-%N-%D %H:%M:%S")

; get domain size
tmp = wrf_user_getvar(a, "pressure", 0)
domainSize = dimsizes(tmp)
ySize = domainSize(1)
xSize = domainSize(2)
print(domainSize)
print(ySize)
print(xSize)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;                   Loop for each station                       ;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

do istn = 0, nstns - 1
    ; check if the domain of station is the opened domain
    if ( stndomain(istn) .ne. domain) then
        continue
    end if

    ; get station location
    res_location = True
    res_location@returnInt = False ; return float
    loc = wrf_user_ll_to_ij(a, stnlon(istn), stnlat(istn), res_location) - 1
    x = loc(0)
    y = loc(1)
    print(loc)

    ; check if the station location is out of the domain
    if ( (y .lt. 0) .or. (y .gt. ySize-1) .or. (x .lt. 0) .or. (x .gt. xSize-1)) then
        print("the station " + stnlist(istn) + " is outside of the domain, skipping.")
        continue
    end if

    output_filename = output_prefix + "/wind_output_" + stnlist(istn) +  "_" + domain +  "_" + start_date + ".txt"
    print("Creating: " + output_filename)
    ;;; latlon to XY ;;;
    res = True
    res@returnInt = False   ; return real values
    loc = wrf_user_ll_to_ij(a, stnlon(istn), stnlat(istn), res) - 1

    sxc0 = floattointeger(loc(0))
    syc0 = floattointeger(loc(1))

    syc1 = syc0 + 1
    sxc1 = sxc0 + 1

    print ("X/Y location is: " + loc)
    yy = (/syc0, syc1/)
    xx = (/sxc0, sxc1/)
    print("X0 is: " + sxc0)
    print("Y0 is: " + syc0)
    print("X1 is: " + sxc1)
    print("Y1 is: " + syc1)

    line_array = new(ntimes(0)+1, "string")
    i_line = 0
    line = "time, "
    do ilvl = 0, nlevel - 1
        line = line + ", " + tostring(levels(ilvl)) + "m windsp" \
                    + ", " + tostring(levels(ilvl)) + "m winddir"
    end do
    line_array(i_line) = line
    i_line = i_line + 1

    ;;; loop for time steps ;;;
    do istep = 0, ntimes(0) - 1

        ;;; read variables
        pres_hpa = wrf_user_getvar(a, "pressure", istep)
        uvmet = wrf_user_getvar(a, "uvmet", istep)
        uvmet10 = wrf_user_getvar(a, "uvmet10", istep)
        umet = uvmet(0, :, :, :)
        vmet = uvmet(1, :, :, :)

        ;;; Load Terrain Data ;;;
        ter = wrf_user_getvar(a,"ter",istep)
        z = wrf_user_getvar(a,"z",istep)
        terrainDims  = dimsizes(ter)
        zdim = dimsizes(z)

;        print(terrainDims)
;        yi = ispan(0, terrainDims(0)-1, 1)
;        xi = ispan(0, terrainDims(1)-1, 1)
;        tmp = linint2(xi, yi, ter, False, loc(1), loc(0), 0)
;        print(tmp)

        ; calculate wind of four corner
        uv_arr0 = wrf_uv_intrp_vert(ter, z, uvmet, uvmet10, syc0, sxc0, levels)
;        uv_arr1 = wrf_uv_intrp_vert(ter, z, uvmet, uvmet10, syc0, sxc1, levels)
;        uv_arr2 = wrf_uv_intrp_vert(ter, z, uvmet, uvmet10, syc1, sxc0, levels)
;        uv_arr3 = wrf_uv_intrp_vert(ter, z, uvmet, uvmet10, syc1, sxc1, levels)

        uv_arr = uv_arr0

;        uv_arr = new(dimsizes(uv_arr0), "float")
;        do ilvl = 0, dimsizes(levels)-1
;            uplane = (/ (/uv_arr0(0, ilvl), uv_arr1(0, ilvl)/), (/uv_arr2(0, ilvl), uv_arr3(0, ilvl)/)/)
;            vplane = (/ (/uv_arr0(1, ilvl), uv_arr1(1, ilvl)/), (/uv_arr2(1, ilvl), uv_arr3(1, ilvl)/)/)
;            uv_arr(0, ilvl) = linint2( xx, yy, uplane, False, loc(0), loc(1), 0)
;            uv_arr(1, ilvl) = linint2( xx, yy, vplane, False, loc(0), loc(1), 0)
;        end do

        spd_arr = sqrt(uv_arr(0,:)^2 + uv_arr(1,:)^2)
        dir_arr = atan2(uv_arr(0,:), uv_arr(1,:)) * r2d + 180

        ;;; write to txt file ;;;
        line = stnnum(istn) + ", " + stnlist(istn) + ", " + timeStr_cst(istep)
        do ilvl = 0, nlevel - 1
            line = line + ", " + \
                sprintf("%9.4f", spd_arr(ilvl)) + ", " + \
                sprintf("%9.4f", dir_arr(ilvl))
        end do
        line_array(i_line) = line
        i_line = i_line + 1
    end do
    asciiwrite(output_filename, line_array)
end do
